

#!/usr/bin/env python


import rospy
from trajectory_msgs.msg import MultiDOFJointTrajectory

import numpy as np
import acado
import time
import matplotlib.pyplot as plt
from numpy import sin, cos,exp
from tf.transformations import euler_from_quaternion, quaternion_from_euler
from ros_numpy import numpify,msgify
from visualization_msgs.msg import MarkerArray,Marker
import time
from trajectory_msgs.msg import MultiDOFJointTrajectory
from sensor_msgs.msg import PointCloud2
from geometry_msgs.msg import Pose,PoseStamped,PoseArray

import open3d as o3d

pc_size=30

N=20 #timestep
NU=2
NOD=pc_size*3+3
NX=4
NY=6
NYN=4

NW=6
NWN=4

pc_W=np.ones([3,pc_size])


def pc2_msg_to_np(pc_msg):
    pc_np=numpify(pc_msg)
    pc_np= np.stack([pc_np[f] for f in ('x', 'y', 'z')])

    return pc_np

def np_to_pc2_msg(pc_np):
    data=np.zeros(pc_np.shape[1],
    dtype=[('x',np.float32), ('y',np.float32), ('z',np.float32)])
    data['x']=pc_np[0,]
    data['y']=pc_np[1,]
    data['z']=pc_np[2,]

    pc_msg=msgify(PointCloud2,data)
    pc_msg.is_dense=False
        
    return pc_msg

def pc_callback(pc_msg):
    #downsample by o3d
    cloud = o3d.geometry.PointCloud()
    pc_np=pc2_msg_to_np(pc_msg)
    cloud.points=o3d.utility.Vector3dVector(pc_np.T)
    
    
    target_size=pc_size-10
    sampling_ratio=float(target_size)/pc_np.shape[1]
    
    downsampled_cloud=cloud.uniform_down_sample(int(1/sampling_ratio))
    pc_downsamped_np=np.asarray(downsampled_cloud.points).T
    np.save('pc_test.npy',pc_downsamped_np)
    pc_downsamped_msg=np_to_pc2_msg(pc_downsamped_np)
    pc_downsamped_msg.header.frame_id="world"
    
    pc_pub.publish(pc_downsamped_msg)
    
    

def multi_dof_trajectory_callback(msg):
    # Process the multi DOF joint trajectory message
    # Extract and print the poses and twists from the message

    new_msg=msg
    
    x0=np.zeros(NX).reshape(1,NX)
    od=np.zeros(NOD)
    y=np.zeros([N,NY])
    u=np.zeros([N,NU])
    Q=np.eye(NW)
    Q[2,2]=10
    Q[3,3]=10
    W=np.transpose(np.tile(Q,N))
    WN=Q[:-2,:-2]
    X=np.tile(x0.reshape(NX,1),N+1).T
    OD=np.ones([N+1,NOD])

    OD[:,pc_size*0:pc_size*1]=pc_W[0,:]  #x
    OD[:,pc_size*1:pc_size*2]=pc_W[1,:]  #y
    OD[:,pc_size*2:pc_size*3]=pc_W[2,:]  #z
    
    xlog=[]
    
    marker_array=MarkerArray()
    pose_array=PoseArray()
    pose_array.header.frame_id="world"
    for p in range(len(msg.points)):
        X,u=acado.mpc(0,1,x0,X,u,OD,y,W,WN,0)
        x0=X[1,:].reshape(1,NX)
        xlog.append(x0)
        
        q=quaternion_from_euler(0,0,x0[0,0])
        new_msg.points[p].transforms[0].rotation.x=q[0]
        new_msg.points[p].transforms[0].rotation.y=q[1]
        new_msg.points[p].transforms[0].rotation.z=q[2]
        new_msg.points[p].transforms[0].rotation.w=q[3]
        
        
        
        #publish marker array
        marker=Marker()
        marker.header.frame_id='world'
        marker.type=Marker.ARROW
        #marker.action=Marker.ADD
        waypoint=new_msg.points[p].transforms[0].translation
        marker.pose.position.x=waypoint.x
        marker.pose.position.y=waypoint.y
        marker.pose.position.z=waypoint.z
        marker.pose.orientation.x=q[0]
        marker.pose.orientation.y=q[1]
        marker.pose.orientation.z=q[2]
        marker.pose.orientation.w=q[3]
        
        pose=marker.pose
        pose_array.poses.append(pose)
        
            # Set the scale, color, and other properties of the marker
        marker.scale.x = 1.0
        marker.scale.y = 0.1
        marker.scale.z = 0.1
        marker.color.a = 1.0
        marker.color.r = 1.0
        marker.color.g = 0.0
        marker.color.b = 0.0

        marker_array.markers.append(marker)
        
        #marker_pub.publish(marker_array)
        pose_pub.publish(pose_array)
    #xlog_arr=np.vstack(xlog)
        

if __name__ == '__main__':
    rospy.init_node('multi_dof_trajectory_subscriber')  # Initialize the ROS node
    rospy.Subscriber('/position_traj', MultiDOFJointTrajectory, multi_dof_trajectory_callback)  # Subscribe to the "MultiDOFJointTrajectory" topic
    rospy.Subscriber('/slot/pointcloud',PointCloud2,pc_callback)
    traj_pub=rospy.Publisher('/active/trajectory',MultiDOFJointTrajectory,queue_size=10)
    pc_pub=rospy.Publisher('/slot/downsampled_pointcloud',PointCloud2,queue_size=10)
    marker_pub=rospy.Publisher('/slot/marker_array',MarkerArray,queue_size=10)
    pose_pub=rospy.Publisher('/slot/pose_array',PoseArray,queue_size=10)
    rospy.spin()  # Keep the node running until it is stopped